from env import ENVIRONMENT  # 从env模块导入ENVIRONMENT类
from Brain import DQN  # 从Brain模块导入DQN类
import numpy as np  # 导入numpy库，用于数组和数学运算
import time  # 导入time库，用于时间相关操作
import random  # 导入random库，用于生成随机数

def main(max_iter):  # 定义主函数，接收最大迭代次数作为参数
    start = time.time()  # 记录开始时间
    agent_reward_list = []  # 初始化记录智能体奖励的列表
    TDMA_reward_list = []  # 初始化记录TDMA节点奖励的列表
    SA_reward_list = []  # 初始化记录SA节点奖励的列表
    reward_list = []  # 初始化记录总奖励的列表
    agent_state_list = []  # 初始化记录智能体状态的列表
    agent_action_list = []  # 初始化记录智能体动作的列表
    state = env.reset()  # 重置环境状态
    counter2 = 0  # 初始化计数器，用于经验回放
    a = 0  # 初始化变量a，具体作用未知

    for i in range(max_iter):  # 进行最大迭代次数的循环
        agent_action = dqn_agent.choose_action(state)  # 智能体选择动作
        # 执行动作并获取观察结果、奖励、智能体奖励、TDMA奖励和SA奖励
        observation_, reward, agent_reward, TDMA_reward, SA_reward = env.step(agent_action)
        #print("SA_reward:",SA_reward)
        # 如果迭代次数小于DRL延迟的两倍减1，则重置环境状态
        if i < env.DRL_delay * 2 - 1:
            next_state = env.reset()
        else:
            # 使用经验回放机制构建下一个状态
            next_state = np.concatenate([agent_state_list[counter2][2:], [agent_action_list[counter2], observation_]])
            counter2 += 1  # 增加计数器

        # 如果迭代次数大于等于DRL延迟的两倍减1，则存储转换到经验回放中
        if i >= env.DRL_delay * 2 - 1:
            dqn_agent.store_transition(state, agent_action, reward, next_state)

        # 更新状态和动作列表
        agent_state_list.append(state)
        agent_action_list.append(agent_action)

        # 更新奖励列表
        agent_reward_list.append(agent_reward)  # 记录智能体奖励
        TDMA_reward_list.append(TDMA_reward)  # 记录TDMA奖励
        SA_reward_list.append(SA_reward)  # 记录SA奖励
        reward_list.append(reward)  # 记录总奖励

        # 省略了自适应训练机制的代码

        # 如果迭代次数超过200且a为0，则开始学习
        if i > 200 and a == 0:
            dqn_agent.learn()

        # 更新状态为下一个状态
        state = next_state

        # 打印当前迭代次数
        print(i)

    # 将智能体奖励、TDMA奖励和SA奖励写入文件
    np.savetxt('rewards/agent_len5e4_M40.txt', agent_reward_list)
    np.savetxt('rewards/TDMA_len5e4_M40.txt', TDMA_reward_list)
    np.savetxt('rewards/SA_len5e4_M40.txt', SA_reward_list)
    # with open('rewards/agent_len5e4_M40.txt', 'w') as my_agent:
    #     for i in agent_reward_list:
    #         my_agent.write(str(i) + '   ')
    # with open('rewards/TDMA_len5e4_M40.txt', 'w') as my_TDMA:
    #     for i in TDMA_reward_list:
    #         my_TDMA.write(str(i) + '   ')
    # with open('rewards/SA_len5e4_M40.txt', 'w') as my_SA:
    #     for i in SA_reward_list:
    #         my_SA.write(str(i) + '   ')

    # 打印平均奖励和运行时间
    #print('-----------------------------')  # 打印分隔符
    # 计算并打印最后2000个时间槽的平均智能体奖励、TDMA奖励、SA奖励和总奖励
    print('average agent reward: {}'.format(np.mean(agent_reward_list[-2000:])))
    print('average TDMA reward:  {}'.format(np.mean(TDMA_reward_list[-2000:])))
    print('average SA reward:  {}'.format(np.mean(SA_reward_list[-2000:])))
    print('average total reward: {}'.format(np.mean(agent_reward_list[-2000:]) + np.mean(TDMA_reward_list[-2000:])+ np.mean(SA_reward_list[-2000:])))
    print('Time elapsed:', time.time() - start)  # 打印程序运行时间

if __name__ == "__main__":  # 如果该脚本作为主程序运行
    # 创建环境实例
    env = ENVIRONMENT(state_size=60)
    # 创建DQN智能体实例，配置智能体参数
    dqn_agent = DQN(env.state_size,  # 智能体的状态数量
                    env.n_actions,  # 智能体的动作数量
                    env.n_nodes,  # 节点数量
                    state_action_memory_size=2*(env.DRL_delay),  # 状态-动作记忆体大小
                    memory_size=500,  # 经验回放缓冲区大小
                    replace_target_iter=200,  # 目标网络参数更新迭代间隔
                    batch_size=32,  # 训练批次大小
                    learning_rate=0.01,  # 学习率
                    gamma=0.95,  # 折扣因子
                    epsilon=0.1,  # 初始探索率
                    epsilon_min=0.001,  # 最小探索率
                    epsilon_decay=0.996,  # 探
                    # 索率衰减率
                    )
    # 调用main函数运行模拟，最大迭代次数为10000
    main(max_iter=10000)